from fastapi import APIRouter

from sparkai.llm.llm import ChatSparkLLM, ChunkPrintHandler
from sparkai.core.messages import ChatMessage
from routers.models import ChatBody

rooter = APIRouter(tags=["饮食生成页面"])

SPARKAI_URL = 'wss://spark-api.xf-yun.com/v1.1/chat'
SPARKAI_APP_ID = 'cd580f0d'
SPARKAI_API_SECRET = 'M2VjMTcwODVhNmViZmM4N2FhYzRjYmRl'
SPARKAI_API_KEY = 'ba412c3aa1073c6066a541da4f7edc39'
SPARKAI_DOMAIN = 'lite'

spark = ChatSparkLLM(
    spark_api_url=SPARKAI_URL,
    spark_app_id=SPARKAI_APP_ID,
    spark_api_key=SPARKAI_API_KEY,
    spark_api_secret=SPARKAI_API_SECRET,
    spark_llm_domain=SPARKAI_DOMAIN,
    streaming=False,
)
handler = ChunkPrintHandler()
spark_ultra = ChatSparkLLM(
    spark_api_url='wss://spark-api.xf-yun.com/v4.0/chat',
    spark_app_id='cd580f0d',
    spark_api_key='ba412c3aa1073c6066a541da4f7edc39',
    spark_api_secret='M2VjMTcwODVhNmViZmM4N2FhYzRjYmRl',
    spark_llm_domain='4.0Ultra',
    streaming=False,
)

@rooter.post("/")
async def check(chatmsg: ChatBody):
    print(chatmsg)

    #将用户信息写入预备知识中
    sys_content="你现在是一个健康饮食和营养饮食方面的专家，请你考虑用户的身体数据做出回答，用户的身体数据如下：性别:"+chatmsg.gender+",年龄:"+str(chatmsg.age)+"岁,体重:"+str(chatmsg.weight)+",身高:"+str(chatmsg.height)
    messages = [ChatMessage(
        role="system",
        content=sys_content
    ),ChatMessage(
        role="user",
        content=chatmsg.body
    )]

    # 根据所选模型进行api调用
    if not chatmsg.type:
        print("调用lite模型")
        a = spark.generate([messages], callbacks=[handler])
    else:
        print("调用4.0Ultra模型")
        a = spark_ultra.generate([messages], callbacks=[handler])
    msg = a.generations[0][0].text

    return msg
